﻿using System;
using System.Collections.Generic;
using System.Drawing;
using System.IO;
using System.Linq;
using System.Text;
using MentalAlchemy.Atomics;

namespace MentalAlchemy.Molecules.MachineLearning
{
	public class MNISTExtractor
	{
		public const int HEIGHT = 28;
		public const int WIDTH = 28;
		public const int PIXEL_COUNT = HEIGHT * WIDTH;

		public const string DEFAULT_TRAIN_LABELS_FOLDER = "train-labels-idx1-ubyte";
		public const string DEFAULT_TRAIN_IMAGES_FOLDER = "train-images-idx3-ubyte";
		public const string DEFAULT_TEST_LABELS_FOLDER = "t10k-labels-idx1-ubyte";
		public const string DEFAULT_TEST_IMAGES_FOLDER = "t10k-images-idx3-ubyte";

		public const string DEFAULT_TRAIN_LABELS_FILE = "train-labels.idx1-ubyte";
		public const string DEFAULT_TRAIN_IMAGES_FILE = "train-images.idx3-ubyte";
		public const string DEFAULT_TEST_LABELS_FILE = "t10k-labels.idx1-ubyte";
		public const string DEFAULT_TEST_IMAGES_FILE = "t10k-images.idx3-ubyte";

		/// <summary>
		/// [molecule]
		/// 
		/// Extracts training and test data, located in default subfolders in the given directory.
		/// </summary>
		/// <param name="folder"></param>
		/// <param name="labelsTrain"></param>
		/// <param name="imgTrain"></param>
		/// <param name="labelsTest"></param>
		/// <param name="imgTest"></param>
		public static void ExtractEverything (string folder, out List<byte> labelsTrain, out List<byte[,]> imgTrain, out List<byte> labelsTest, out List<byte[,]> imgTest)
		{
			var path = Path.Combine(folder, DEFAULT_TRAIN_LABELS_FOLDER);
			path = Path.Combine(path, DEFAULT_TRAIN_LABELS_FILE);
			labelsTrain = GetLabels(path);

			path = Path.Combine(folder, DEFAULT_TRAIN_IMAGES_FOLDER);
			path = Path.Combine(path, DEFAULT_TRAIN_IMAGES_FILE);
			imgTrain = GetImages(path);

			path = Path.Combine(folder, DEFAULT_TEST_LABELS_FOLDER);
			path = Path.Combine(path, DEFAULT_TEST_LABELS_FILE);
			labelsTest = GetLabels(path);

			path = Path.Combine(folder, DEFAULT_TEST_IMAGES_FOLDER);
			path = Path.Combine(path, DEFAULT_TEST_IMAGES_FILE);
			imgTest = GetImages(path);
		}

		/// <summary>
		/// [molecule]
		/// 
		/// Extracts labels from the given file.
		/// </summary>
		/// <param name="filename"></param>
		/// <returns></returns>
		public static List<byte> GetLabels(string filename)
		{
			List<byte> res;

			using (var reader = new BinaryReader(File.Open(filename, FileMode.Open)))
			{
				reader.ReadInt32();	// skip magic number.
				var count = (uint)reader.ReadInt32();
				count = Utils.ReverseBytesOrder(count);
				res = new List<byte>(new byte[count]);

				for (int i = 0; i < count; i++)
				{
					res[i] = reader.ReadByte();
				}
			}

			return res;
		}

		/// <summary>
		/// [molecule]
		/// 
		/// Extracts labels from the given file.
		/// </summary>
		/// <param name="filename"></param>
		/// <returns></returns>
		public static List<byte[,]> GetImages(string filename)
		{
			var res = new List<byte[,]>();

			using (var reader = new BinaryReader(File.Open(filename, FileMode.Open)))
			{
				reader.ReadInt32();	// skip magic number.
				var count = (uint)reader.ReadInt32();
				count = Utils.ReverseBytesOrder(count);
				reader.ReadInt32();	// skip dimensions.
				reader.ReadInt32();	// skip dimensions.

				for (int i = 0; i < count; i++)
				{
					var bytes = reader.ReadBytes(PIXEL_COUNT);
					res.Add(MatrixMath.CreateFromVector(bytes, WIDTH));
				}
			}

			return res;
		}
	}
}
